from .dataset import FundusDataset
from torch.utils.data import DataLoader
from torchvision import transforms
import lightning as L

class FundusDataModule(L.LightningDataModule):
    def __init__(self, data_dir, batch_size=32, input_size=224,num_workers=4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.input_size = input_size
        self.mean = (0.4237, 0.2609, 0.1284)
        self.std = (0.2948, 0.2016, 0.1366) # from retfound
        self.num_workers = num_workers
    def on_train_start(self):
        activate_device = self.device  # 获取当前设备
        self.logger.info(f"Using device {activate_device}")  # 记录到 logger
        print(f"Using device {activate_device}")  # 控制台打
    def setup(self, stage=None):
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(self.input_size),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.std),
        ])

        test_transform = transforms.Compose([
            transforms.Resize((self.input_size, self.input_size)),
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.std),
        ])

        self.train_dataset = FundusDataset(root=self.data_dir, split='train', transform=train_transform)
        self.test_dataset = FundusDataset(root=self.data_dir, split='test', transform=test_transform)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers)
